
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '0'

from diffusers import FluxKontextPipeline, FluxControlPipeline, FluxPriorReduxPipeline
from util_flux import process_img_1024, vertical_concat_images, horizontal_concat_images, resize_with_aspect
from PIL import Image
from itertools import product
from util_for_os import osj, ose
from MODEL_CKP import FLUX_KONTEXT, FLUX_REDUX
import torch

def get_mask(size=(1024, 1024), fill=(255, 255, 255)):
    """
    生成指定大小和颜色的纯色掩码图像
    """
    return Image.new('RGB', size, color=fill)

def collect_image_paths(directory, num=5, suffix='.jpg'):
    """
    从指定目录收集前num个指定后缀的图片路径
    """
    image_list = []
    for entry in os.scandir(directory):
        filename = entry.name
        if not filename.endswith(suffix):
            continue
        if len(image_list) < num:
            image_list.append(osj(directory, filename))
        else:
            break
    return image_list

def get_prompt_embeds(pipe_prior_redux, control_image, redux_image, alpha=0.8):
    """
    获取并混合两个图片的prompt_embeds和pooled_prompt_embeds
    """
    with torch.no_grad():
        prompt_emb, pooled_prompt_emb = pipe_prior_redux(control_image, return_dict=False)
        prompt_emb2, pooled_prompt_emb2 = pipe_prior_redux(redux_image, return_dict=False)
        # 混合两个图片的embeddings
        prompt_emb_mix = (1 - alpha) * prompt_emb + alpha * prompt_emb2
        pooled_prompt_emb_mix = (1 - alpha) * pooled_prompt_emb + alpha * pooled_prompt_emb2
    return prompt_emb2, pooled_prompt_emb2  # 注意：原代码实际只用prompt_emb2, pooled_prompt_emb2

def generate_and_save_grid(types, num=5, save_prefix='tmp_redux'):
    """
    针对每个服饰类型，生成交叉组合的图片，并保存拼接结果
    """
    # 初始化pipeline
    pipe = FluxKontextPipeline.from_pretrained(FLUX_KONTEXT, torch_dtype=torch.bfloat16)
    pipe.to("cuda")
    pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(FLUX_REDUX, torch_dtype=torch.bfloat16).to("cuda")

    examples_dir = lambda t: f'/mnt/nas/shengjie/datasets/cloth_{t}_localimg'

    for t_id, t in enumerate(types):
        # 收集图片路径
        dir1, dir2 = examples_dir(t), examples_dir(t)
        clo_list1 = collect_image_paths(dir1, num)
        clo_list2 = collect_image_paths(dir2, num)

        # 构建结果图片表格的初始行和列
        res_imgs = [
            [get_mask()] + [process_img_1024(path) for path in clo_list1],  # 第一行：空白+clo_list1图片
            *[[process_img_1024(path)] for path in clo_list2],              # 后续行：clo_list2图片
        ]

        idx = 0
        for c1, c2 in product(clo_list1, clo_list2):
            """
            构建如下表格：
            /        原图 c1-1  c1-2  c1-3 ...
            原图c2-1     合成1  合成2 合成3 ...
            c2-2        ...
            c2-3
            ...
            """
            control_image_path = c1
            redux_image_path = c2

            control_image = process_img_1024(control_image_path)
            redux_image = process_img_1024(redux_image_path)

            # 获取prompt_embeds和pooled_prompt_embeds
            prompt_emb2, pooled_prompt_emb2 = get_prompt_embeds(pipe_prior_redux, control_image, redux_image, alpha=0.8)

            # 生成新图片
            with torch.no_grad():
                image = pipe(
                    image=control_image,
                    height=control_image.height,
                    width=control_image.width,
                    num_inference_steps=8,
                    guidance_scale=2.5,
                    prompt_embeds=prompt_emb2,
                    pooled_prompt_embeds=pooled_prompt_emb2,
                ).images[0]

            # 将生成的图片插入到对应的表格位置
            res_imgs[idx % num + 1].append(process_img_1024('', img_pil=image))
            idx += 1

            # 保存每次生成的横向拼接结果（可选，调试用）
            concat_tmp_res = horizontal_concat_images([
                process_img_1024(control_image_path),
                redux_image,
                process_img_1024('', img_pil=image),
            ])
            concat_tmp_res.save('tmp_redux2.jpg')

        # 拼接整张表格图片并保存
        res_imgs_hori_concat = [horizontal_concat_images(ri) for ri in res_imgs]
        res_imgs_verti_concat = vertical_concat_images(res_imgs_hori_concat)
        res_imgs_verti_concat.save(f'{save_prefix}_{t_id}.jpg')

if __name__ == "__main__":
    """
    主程序入口：对每种服饰类型生成交叉组合图片表格
    """
    types = ['collar', 'sleeve', 'pockets']
    generate_and_save_grid(types, num=5, save_prefix='tmp_redux')